package com.alibaba.nacos.core.remote.grpc;

import com.alibaba.nacos.api.utils.StringUtils;
import com.alibaba.nacos.core.remote.ConnectionManager;
import io.grpc.Attributes;
import io.grpc.Grpc;
import io.grpc.ServerTransportFilter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import java.net.InetSocketAddress;

@Slf4j
@RequiredArgsConstructor
public class AddressTransportFilter extends ServerTransportFilter {

    private final ConnectionManager connectionManager;

    @Override
    public Attributes transportReady(Attributes transportAttrs) {
        InetSocketAddress remoteAddress = (InetSocketAddress) transportAttrs.get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR);
        InetSocketAddress localAddress = (InetSocketAddress) transportAttrs.get(Grpc.TRANSPORT_ATTR_LOCAL_ADDR);
        int remotePort = remoteAddress.getPort();
        int localPort = localAddress.getPort();
        String remoteIp = remoteAddress.getAddress().getHostAddress();

        Attributes attrWrapper = transportAttrs.toBuilder()
                .set(GrpcServerConstants.ATTR_TRANS_KEY_CONN_ID, System.currentTimeMillis() + "_" + remoteIp + "_" + remotePort)
                .set(GrpcServerConstants.ATTR_TRANS_KEY_REMOTE_IP, remoteIp)
                .set(GrpcServerConstants.ATTR_TRANS_KEY_REMOTE_PORT, remotePort)
                .set(GrpcServerConstants.ATTR_TRANS_KEY_LOCAL_PORT, localPort)
                .build();

        return attrWrapper;
    }

    @Override
    public void transportTerminated(Attributes transportAttrs) {
        if (transportAttrs == null) {
            return;
        }

        String connectionId = transportAttrs.get(GrpcServerConstants.ATTR_TRANS_KEY_CONN_ID);
        if (StringUtils.isBlank(connectionId)) {
            return;
        }

        log.info("connectionId:{}", connectionId);
        connectionManager.unregister(connectionId);
    }
}
